{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "86a13ed9",
   "metadata": {},
   "source": [
    "Creating a submodel\n",
    "\n",
    "In [Tutorial 3](./3-negative-particle-problem.ipynb) we showed how to create a simple model from scratch in PyBaMM. In this tutorial we will solve a very similar problem, but using separate submodels for the linear diffusion problem and the model for the surface flux. In this simple example the surface flux is just some known function of the concentration, so we could just explicitly define it in the model for diffusion. However, we write it as a separate model to show how submodels interact. \n",
    "\n",
    "We solved the problem of linear diffusion on a unit sphere with a flux at the boundary that depends on the concentration\n",
    "$$\n",
    "  \\frac{\\partial c}{\\partial t} = \\nabla \\cdot (\\nabla c),\n",
    "$$\n",
    "with the following boundary and initial conditions:\n",
    "$$\n",
    "  \\left.\\frac{\\partial c}{\\partial r}\\right\\vert_{r=0} = 0, \\quad \\left.\\frac{\\partial c}{\\partial r}\\right\\vert_{r=1} = -j, \\quad \\left.c\\right\\vert_{t=0} = c_0,\n",
    "$$\n",
    "where\n",
    "$$\n",
    "j = \\left.j_0(1-c)^{1/2}c^{1/2}\\right\\vert_{r=1}\n",
    "$$\n",
    "Here $c_0$ and $j_0$ are parameters we can control. Again we will assume that everything is non-dimensional and focus on how to set up and solve the model rather than any specific physical interpretation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "9e8cf744",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install \"pybamm[plot,cite]\" -q    # install PyBaMM if it is not installed\n",
    "import pybamm"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "2d074d41",
   "metadata": {},
   "source": [
    "## Setting up the model\n",
    "\n",
    "Again we start with an empty `BaseModel` class"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "1dff6cd1",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = pybamm.BaseModel()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "c44d3484",
   "metadata": {},
   "source": [
    "Next we set up the submodel for diffusion in the particle using a `pybamm.BaseSubModel`. Each submodel has methods that define the variables, equations, initial and boundary conditions corresponding to the physics the model describes.\n",
    "\n",
    "First `get_fundamental_variables` defines any variables that can be defined independently of any other submodels that may be included in our model. Here we can define the concentration, surface concentration and flux, and add them to the dictionary of variables. \n",
    "\n",
    "Next we can use `get_coupled_variables` to define any variables that _do_ depend on variables defined in another submodel. In this simple example we don't have any variables to define here. However, if we had included a temperature dependent diffusivity, for example, then we would have needed to define the flux in `get_coupled_variables` since it would now depend on the temperature which would be defined in a separate submodel.\n",
    "\n",
    "Once we have defined all the variables we need we can write down our equations. Any equations that include time derivatives will turn into ordinary differential equations after discretisation. We set the right hand sides of such equations in the `set_rhs` method. In this example we add the right hand side of the diffusion equation to the `rhs` dictionary. Equations that don't contain time derivatives give algebraic constraints in our model. These equations are set in the `set_algebraic` method. In this example we don't have any algebraic equations, so we can skip this method.\n",
    "\n",
    "Finally we set the boundary and initial conditions using the methods `set_boundary_conditions` and `set_initial_conditions`, respectively. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "337c1475",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Particle(pybamm.BaseSubModel):\n",
    "    def __init__(self, param, domain, options=None):\n",
    "        super().__init__(param, domain, options=options)\n",
    "\n",
    "    def get_fundamental_variables(self):\n",
    "        # create concentration variable\n",
    "        c = pybamm.Variable(\"Concentration\", domain=\"negative particle\")\n",
    "\n",
    "        # define concentration at the surface of the sphere\n",
    "        c_surf = pybamm.surf(c)\n",
    "\n",
    "        # define flux\n",
    "        N = -pybamm.grad(c)\n",
    "\n",
    "        # create dictionary of model variables\n",
    "        variables = {\n",
    "            \"Concentration\": c,\n",
    "            \"Surface concentration\": c_surf,\n",
    "            \"Flux\": N,\n",
    "        }\n",
    "\n",
    "        return variables\n",
    "\n",
    "    def get_coupled_variables(self, variables):\n",
    "        return variables\n",
    "\n",
    "    def set_rhs(self, variables):\n",
    "        # extract the variables we need\n",
    "        c = variables[\"Concentration\"]\n",
    "        N = variables[\"Flux\"]\n",
    "\n",
    "        # define the rhs of the PDE\n",
    "        dcdt = -pybamm.div(N)\n",
    "\n",
    "        # add it to the submodel dictionary\n",
    "        self.rhs = {c: dcdt}\n",
    "\n",
    "    def set_algebraic(self, variables):\n",
    "        pass\n",
    "\n",
    "    def set_boundary_conditions(self, variables):\n",
    "        # extract the variables we need\n",
    "        c = variables[\"Concentration\"]\n",
    "        j = variables[\"Boundary flux\"]\n",
    "\n",
    "        # add the boundary conditions to the submodel dictionary\n",
    "        self.boundary_conditions = {\n",
    "            c: {\"left\": (0, \"Neumann\"), \"right\": (-j, \"Neumann\")}\n",
    "        }\n",
    "\n",
    "    def set_initial_conditions(self, variables):\n",
    "        # extract the variable we need\n",
    "        c = variables[\"Concentration\"]\n",
    "\n",
    "        # define the initial concentration parameter\n",
    "        c0 = pybamm.Parameter(\"Initial concentration\")\n",
    "\n",
    "        # add the initial conditions to the submodel dictionary\n",
    "        self.initial_conditions = {c: c0}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "642bbd46",
   "metadata": {},
   "outputs": [],
   "source": [
    "class BoundaryFlux(pybamm.BaseSubModel):\n",
    "    def __init__(self, param, domain, options=None):\n",
    "        super().__init__(param, domain, options=options)\n",
    "\n",
    "    def get_coupled_variables(self, variables):\n",
    "        # extract the variable we need\n",
    "        c_surf = variables[\"Surface concentration\"]\n",
    "\n",
    "        # define the flux parameter\n",
    "        j0 = pybamm.Parameter(\"Flux parameter\")\n",
    "        j = j0 * (1 - c_surf) ** (1 / 2) * c_surf ** (1 / 2)  # prescribed boundary flux\n",
    "\n",
    "        # update dictionary of model variables\n",
    "        variables.update({\"Boundary flux\": j})\n",
    "\n",
    "        return variables"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "62907bc5",
   "metadata": {},
   "source": [
    "We can now set the submodels in a model by assigning a dictionary to `model.submodels`. The dictionary key is the name we want to give to the submodel and the value is an instance of the submodel class we want to use.\n",
    "\n",
    "When we instantiate a submodel we are required to pass in `param`, a class of parameter symbols we are going to call, and `domain`, the domain on which the submodel lives. In this example we will simply set `param` to `None` and hard-code the definition of our parameters into the submodel. When writing lots of submodels it is more efficient to define _all_ the parameters in a shared class, and pass this to each submodel. For the domain we will choose \"Negative\". "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "7fed4b8a",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.submodels = {\n",
    "    \"Particle\": Particle(None, \"Negative\"),\n",
    "    \"Boundary flux\": BoundaryFlux(None, \"Negative\"),\n",
    "}"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "9e910625",
   "metadata": {},
   "source": [
    "At this stage we have just told the model which submodels it is constructed from, but the variables and equations have not yet been created. For example if we look at the `rhs` dictionary it is empty."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "93541aa0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{}"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.rhs"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "c6e07559",
   "metadata": {},
   "source": [
    "To populate the model variables, equations, boundary and initial conditions we need to \"build\" the model. To do this we call `build_model`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "1a8b9548",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.build_model()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "649c0260",
   "metadata": {},
   "source": [
    "This loops through all of the submodels, first creating the \"fundamental variables\", followed by the \"coupled variables\" and finally the equations (`rhs` and `algebraic`) and the boundary and initial conditions. Now we see that `model.rhs` contains our diffusion equation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "374fae5c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{Variable(0x48a91fa41dea72d3, Concentration, children=[], domains={'primary': ['negative particle']}): Divergence(0x4befa5e7cf20d0c5, div, children=['grad(Concentration)'], domains={'primary': ['negative particle']})}"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.rhs"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "a7cab3ac",
   "metadata": {},
   "source": [
    "## Using the model\n",
    "\n",
    "We can now use our model as in the previous tutorial.\n",
    "\n",
    "We first set up our geometry and mesh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "2f2da3bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "r = pybamm.SpatialVariable(\n",
    "    \"r\", domain=[\"negative particle\"], coord_sys=\"spherical polar\"\n",
    ")\n",
    "geometry = {\"negative particle\": {r: {\"min\": 0, \"max\": 1}}}\n",
    "spatial_methods = {\"negative particle\": pybamm.FiniteVolume()}\n",
    "submesh_types = {\"negative particle\": pybamm.Uniform1DSubMesh}\n",
    "var_pts = {r: 20}\n",
    "mesh = pybamm.Mesh(geometry, submesh_types, var_pts)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "2a89e132",
   "metadata": {},
   "source": [
    "and then set up our simulation, remembering to set values for our parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "7589648c",
   "metadata": {},
   "outputs": [],
   "source": [
    "parameter_values = pybamm.ParameterValues(\n",
    "    {\n",
    "        \"Initial concentration\": 0.9,\n",
    "        \"Flux parameter\": 0.8,\n",
    "    }\n",
    ")\n",
    "\n",
    "solver = pybamm.ScipySolver()\n",
    "\n",
    "sim = pybamm.Simulation(\n",
    "    model,\n",
    "    geometry=geometry,\n",
    "    parameter_values=parameter_values,\n",
    "    submesh_types=submesh_types,\n",
    "    var_pts=var_pts,\n",
    "    spatial_methods=spatial_methods,\n",
    "    solver=solver,\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "1cec6e66",
   "metadata": {},
   "source": [
    "Finally we can solve the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "cc9bdc87",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<pybamm.solvers.solution.Solution at 0x7f9dc1b2e490>"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sim.solve([0, 1])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "e6bafd61",
   "metadata": {},
   "source": [
    "and plot the results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "f549d023",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2022-07-11 13:20:04.194 - [WARNING] processed_variable.get_spatial_scale(521): No length scale set for negative particle. Using default of 1 [m].\n",
      "2022-07-11 13:20:04.209 - [WARNING] processed_variable.get_spatial_scale(521): No length scale set for negative particle. Using default of 1 [m].\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d9780b7718bd4a9f932f5dd92c065756",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "interactive(children=(FloatSlider(value=0.0, description='t', max=1.0, step=0.01), Output()), _dom_classes=('w…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<pybamm.plotting.quick_plot.QuickPlot at 0x7f9dc1b2ec10>"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# pass in a list of the variables we want to plot\n",
    "sim.plot([\"Concentration\", \"Surface concentration\", \"Flux\", \"Boundary flux\"])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "3d88c346",
   "metadata": {},
   "source": [
    "In this notebook we saw how to split a model up into submodels. Although this was a simple example it let us understand how to construct submodels and see how they interact via coupled variables."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "f20b7969",
   "metadata": {},
   "source": [
    "## References\n",
    "\n",
    "The relevant papers for this notebook are:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "6301f16b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1] Joel A. E. Andersson, Joris Gillis, Greg Horn, James B. Rawlings, and Moritz Diehl. CasADi – A software framework for nonlinear optimization and optimal control. Mathematical Programming Computation, 11(1):1–36, 2019. doi:10.1007/s12532-018-0139-4.\n",
      "[2] Charles R. Harris, K. Jarrod Millman, Stéfan J. van der Walt, Ralf Gommers, Pauli Virtanen, David Cournapeau, Eric Wieser, Julian Taylor, Sebastian Berg, Nathaniel J. Smith, and others. Array programming with NumPy. Nature, 585(7825):357–362, 2020. doi:10.1038/s41586-020-2649-2.\n",
      "[3] Valentin Sulzer, Scott G. Marquis, Robert Timms, Martin Robinson, and S. Jon Chapman. Python Battery Mathematical Modelling (PyBaMM). Journal of Open Research Software, 9(1):14, 2021. doi:10.5334/jors.309.\n",
      "[4] Pauli Virtanen, Ralf Gommers, Travis E. Oliphant, Matt Haberland, Tyler Reddy, David Cournapeau, Evgeni Burovski, Pearu Peterson, Warren Weckesser, Jonathan Bright, and others. SciPy 1.0: fundamental algorithms for scientific computing in Python. Nature Methods, 17(3):261–272, 2020. doi:10.1038/s41592-019-0686-2.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "pybamm.print_citations()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
